package cn.doitedu.ml.doit13.bayes

import cn.doitedu.commons.util.SparkUtil
import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.types.DataTypes

import scala.collection.mutable

/**
 * 调用sparkmllib中现成的朴素贝叶斯算法，来进行出轨预测（分类）
 */
object NaiveBayesChugui {

  def main(args: Array[String]): Unit = {

    val spark = SparkUtil.getSparkSession(this.getClass.getSimpleName)
    import org.apache.spark.sql.functions._
    import spark.implicits._

    /**
     * 1.  模型训练
     */
    // 加载训练数据集
    val sample = spark.read.option("header", true).csv("userprofile/data/demo/bayes/chugui/sample/sample.csv")

    val arr2vec: UserDefinedFunction = udf((arr: mutable.WrappedArray[String]) => {
      // 用hash映射来生成特征向量 [老师，中，青年，男]
      val features = new Array[Double](20)
      arr.foreach(s=>{
        val idx = (s.hashCode & Integer.MAX_VALUE) % features.size
        features(idx) = 1.0
      })
      Vectors.dense(features)
    })
    spark.udf.register("arr2vec",arr2vec)

    // 将特征 向量化
    val sampleVecs = sample.selectExpr("name",
      "case label when '出轨' then 0.0 else 1.0 end as label",
      "arr2vec(array(job,income,age,sex)) as vec"
    )

    // 构造朴素贝叶斯算法工具
    val bayes = new NaiveBayes()
      .setFeaturesCol("vec")
      .setLabelCol("label")
      .setSmoothing(1.0) // 拉普拉斯平滑系数

    // 用算法工具对训练集训练模型
    val model = bayes.fit(sampleVecs)
    // 保存训练好的模型
    //model.save("userprofile/data/demo/bayes/chugui/model3")


    /**
     *   2.  利用模型做预测
     */
    // 加载之前训练好的模型
    //val model1 = NaiveBayesModel.load("userprofile/data/demo/bayes/chugui/model")
    // 加载待预测数据集
    val test = spark.read.option("header",true).csv("userprofile/data/demo/bayes/chugui/test/test.csv")
    // 对待预测数据集进行特征处理（特征工程）
    test.createTempView("test")
    var testVecs = spark.sql(
      """
        |
        |select
        |name,
        |arr2vec(
        |  array( job,income,age,sex )
        |) as vec
        |from test
        |
        |""".stripMargin)

    // 用模型来预测
    val res = model.transform(testVecs)
    res.show(100,false)

    spark.close()
  }

}
